package ai.example.langchain4j.utils;

import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.Content;
import dev.langchain4j.data.message.ImageContent;
import dev.langchain4j.data.message.TextContent;
import dev.langchain4j.data.message.UserMessage;
import org.springframework.web.multipart.MultipartFile;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;

public class ChatMessageBuilder {

  public static ChatMessage buildChatMessage(String query, List<MultipartFile> files) {
    List<Content> contents = new ArrayList<>();
    contents.add(new TextContent(query));
    appendImageContentsIfExist(files, contents);
    UserMessage userMessage = new UserMessage(contents);
    return userMessage;
  }

  private static void appendImageContentsIfExist(List<MultipartFile> imageFiles, List<Content> contents) {
    if (imageFiles == null || imageFiles.isEmpty()) {
      return;
    }
    imageFiles.stream()
      .filter(file -> file.getContentType().startsWith("image/"))
      .forEach(file -> {
        String mimeType = file.getContentType().split("/")[1]; // png, jpg, jpeg 等
        byte[] fileBytes;
        try {
          fileBytes = file.getBytes();
        } catch (IOException e) {
          throw new RuntimeException(e);
        }
        String base64Data = Base64.getEncoder().encodeToString(fileBytes);
        contents.add(new ImageContent(base64Data, mimeType));
      });
  }

}
